"""
摄像头坐姿检测主程序
"""

import cv2
import numpy as np
import time
from typing import Optional
import argparse
import sys
import os

# 添加当前目录到路径
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

from camera_handler import CameraHandler
from posture_detector import PostureDetector
from config import DISPLAY_CONFIG, COLORS
from utils import draw_text_with_background, get_grade_color, create_info_panel

class PostureDetectionApp:
    """坐姿检测应用程序"""
    
    def __init__(self, camera_id: int = 0, show_debug: bool = False, user_id: str = "default"):
        """
        初始化应用程序
        
        Args:
            camera_id: 摄像头ID
            show_debug: 是否显示调试信息
            user_id: 用户ID，用于个性化学习
        """
        self.camera_handler = CameraHandler(camera_id)
        self.posture_detector = PostureDetector()
        # 使用个性化分析器
        from posture_analyzer import PostureAnalyzer
        self.posture_detector.posture_analyzer = PostureAnalyzer(user_id)
        self.show_debug = show_debug
        self.user_id = user_id
        
        # 统计信息
        self.frame_count = 0
        self.fps_counter = 0
        self.last_fps_time = time.time()
        self.current_fps = 0
        
        # 检测状态
        self.is_running = False
        self.last_detection_time = time.time()
        
        print("摄像头坐姿检测系统已初始化")
        print("精细化评估算法包含以下维度：")
        print("1. 头部与桌面距离 (20%)")
        print("2. 脊柱对齐评估 (25%)")
        print("3. 颈椎姿态评估 (30%)")
        print("4. 肩膀姿态评估 (20%)")
        print("5. 坐姿稳定性评估 (5%)")
    
    def run(self):
        """运行坐姿检测"""
        if not self.camera_handler.open_camera():
            print("摄像头初始化失败")
            return
        
        self.is_running = True
        print("开始坐姿检测...")
        print("按 'q' 键退出，按 'd' 键切换调试模式")
        
        try:
            while self.is_running:
                success, frame = self.camera_handler.read_frame()
                
                if not success:
                    print("读取摄像头帧失败")
                    break
                
                # 处理帧
                processed_frame = self._process_frame(frame)
                
                # 显示结果
                cv2.imshow('智能坐姿检测 - 精细化评估', processed_frame)
                
                # 处理按键
                key = cv2.waitKey(1) & 0xFF
                if key == ord('q'):
                    break
                elif key == ord('d'):
                    self.show_debug = not self.show_debug
                    print(f"调试模式: {'开启' if self.show_debug else '关闭'}")
                elif key == ord('r'):
                    self.posture_detector.reset_detector()
                    print("检测器已重置")
                
                # 更新统计信息
                self._update_statistics()
                
        except KeyboardInterrupt:
            print("用户中断程序")
        except Exception as e:
            print(f"运行时错误: {e}")
        finally:
            self._cleanup()
    
    def _process_frame(self, frame: np.ndarray) -> np.ndarray:
        """
        处理单帧图像
        
        Args:
            frame: 输入帧
        
        Returns:
            处理后的帧
        """
        # 获取坐姿检测结果
        result = self.posture_detector.detect_posture(frame)
        
        # 在帧上绘制结果
        output_frame = self._draw_results(frame, result)
        
        # 如果启用调试模式，显示额外信息
        if self.show_debug:
            output_frame = self._draw_debug_info(output_frame, result)
        
        return output_frame
    
    def _draw_results(self, frame: np.ndarray, result: dict) -> np.ndarray:
        """
        在帧上绘制检测结果
        
        Args:
            frame: 输入帧
            result: 检测结果
        
        Returns:
            绘制后的帧
        """
        output_frame = frame.copy()
        
        if not result['success']:
            # 如果检测失败，显示提示信息
            self._draw_no_detection_message(output_frame)
            return output_frame
        
        metrics = result['metrics']
        landmarks = result['landmarks']
        
        # 绘制关键点和骨架
        if DISPLAY_CONFIG['show_landmarks'] and landmarks:
            self._draw_pose_landmarks(output_frame, landmarks)
        
        # 绘制主要评分信息
        self._draw_main_score(output_frame, metrics)
        
        # 绘制详细指标
        if DISPLAY_CONFIG['show_metrics']:
            self._draw_detailed_metrics(output_frame, metrics)
        
        # 绘制建议
        if DISPLAY_CONFIG['show_suggestions'] and metrics.suggestions:
            self._draw_suggestions(output_frame, metrics.suggestions)
        
        return output_frame
    
    def _draw_pose_landmarks(self, frame: np.ndarray, landmarks: list):
        """绘制人体关键点"""
        # 绘制关键点
        for point in landmarks:
            if point[0] > 0 and point[1] > 0:
                cv2.circle(frame, (int(point[0]), int(point[1])), 
                          DISPLAY_CONFIG['circle_radius'], 
                          COLORS['landmarks'], -1)
        
        # 绘制连接线（简化的骨架）
        connections = [
            (7, 8),   # 左耳-右耳
            (11, 12), # 左肩-右肩
            (11, 13), # 左肩-左肘
            (12, 14), # 右肩-右肘
            (11, 23), # 左肩-左髋
            (12, 24), # 右肩-右髋
            (23, 24), # 左髋-右髋
        ]
        
        for start_idx, end_idx in connections:
            if (start_idx < len(landmarks) and end_idx < len(landmarks) and
                landmarks[start_idx][0] > 0 and landmarks[start_idx][1] > 0 and
                landmarks[end_idx][0] > 0 and landmarks[end_idx][1] > 0):
                
                start_point = (int(landmarks[start_idx][0]), int(landmarks[start_idx][1]))
                end_point = (int(landmarks[end_idx][0]), int(landmarks[end_idx][1]))
                cv2.line(frame, start_point, end_point, 
                        COLORS['skeleton'], DISPLAY_CONFIG['line_thickness'])
    
    def _draw_main_score(self, frame: np.ndarray, metrics):
        """绘制主要评分信息"""
        # 评分背景
        score_color = get_grade_color(metrics.grade)
        
        # 主要评分
        score_text = f"坐姿评分: {metrics.overall_score:.1f}/100"
        draw_text_with_background(
            frame, score_text, (10, 30),
            font_scale=0.8, font_thickness=2,
            text_color=(255, 255, 255), bg_color=score_color
        )
        
        # 等级显示
        grade_names = {
            'excellent': '优秀',
            'good': '良好', 
            'fair': '一般',
            'poor': '较差'
        }
        grade_text = f"等级: {grade_names.get(metrics.grade, '未知')}"
        draw_text_with_background(
            frame, grade_text, (10, 65),
            font_scale=0.6, font_thickness=2,
            text_color=(255, 255, 255), bg_color=score_color
        )
    
    def _draw_detailed_metrics(self, frame: np.ndarray, metrics):
        """绘制详细指标"""
        y_offset = 100
        line_height = 25
        
        # 详细指标列表
        detailed_info = [
            f"头部距离: {metrics.head_desk_distance:.1f}",
            f"脊柱对齐: {metrics.spine_alignment_score:.1f}",
            f"颈椎前倾: {metrics.neck_forward_angle:.1f}°",
            f"颈椎侧倾: {metrics.neck_side_tilt:.1f}°",
            f"肩膀高差: {metrics.shoulder_height_diff:.1f}°",
            f"肩膀对称: {metrics.shoulder_symmetry:.1f}",
            f"坐姿稳定: {metrics.sitting_stability:.1f}"
        ]
        
        for i, info in enumerate(detailed_info):
            draw_text_with_background(
                frame, info, (10, y_offset + i * line_height),
                font_scale=0.5, font_thickness=1,
                text_color=(255, 255, 255), bg_color=(0, 0, 0)
            )
    
    def _draw_suggestions(self, frame: np.ndarray, suggestions: list):
        """绘制改进建议"""
        if not suggestions:
            return
        
        # 建议区域
        suggestion_y = 400
        max_suggestions = 3
        
        # 标题
        draw_text_with_background(
            frame, "改进建议:", (10, suggestion_y),
            font_scale=0.6, font_thickness=2,
            text_color=(255, 255, 255), bg_color=(0, 100, 200)
        )
        
        # 建议内容
        for i, suggestion in enumerate(suggestions[:max_suggestions]):
            # 限制建议长度
            if len(suggestion) > 35:
                suggestion = suggestion[:32] + "..."
            
            draw_text_with_background(
                frame, f"• {suggestion}", (10, suggestion_y + 30 + i * 25),
                font_scale=0.5, font_thickness=1,
                text_color=(255, 255, 255), bg_color=(0, 50, 100)
            )
    
    def _draw_debug_info(self, frame: np.ndarray, result: dict) -> np.ndarray:
        """绘制调试信息"""
        if not result['success']:
            return frame
        
        metrics = result['metrics']
        
        # 右侧调试信息
        debug_x = frame.shape[1] - 250
        debug_y = 30
        line_height = 20
        
        debug_info = [
            f"FPS: {self.current_fps:.1f}",
            f"帧数: {self.frame_count}",
            f"检测耗时: {result.get('processing_time', 0):.3f}s",
            "",
            "详细数值:",
            f"颈椎角度: {metrics.cervical_angle:.1f}°",
            f"胸椎角度: {metrics.thoracic_angle:.1f}°",
            f"脖子扭转: {metrics.neck_twist_ratio:.3f}",
            f"肩膀前后差: {metrics.shoulder_front_back_diff:.1f}",
            f"肩膀内收: {metrics.shoulder_narrowing:.3f}",
        ]
        
        for i, info in enumerate(debug_info):
            if info:  # 跳过空行
                draw_text_with_background(
                    frame, info, (debug_x, debug_y + i * line_height),
                    font_scale=0.4, font_thickness=1,
                    text_color=(255, 255, 255), bg_color=(50, 50, 50)
                )
        
        return frame
    
    def _draw_no_detection_message(self, frame: np.ndarray):
        """绘制无检测结果的提示信息"""
        message = "未检测到人体姿态"
        frame_center = (frame.shape[1] // 2 - 100, frame.shape[0] // 2)
        
        draw_text_with_background(
            frame, message, frame_center,
            font_scale=0.8, font_thickness=2,
            text_color=(255, 255, 255), bg_color=(0, 0, 255)
        )
        
        # 提示信息
        tips = [
            "请确保：",
            "1. 面向摄像头坐好",
            "2. 光线充足",
            "3. 上半身完整可见"
        ]
        
        for i, tip in enumerate(tips):
            draw_text_with_background(
                frame, tip, (frame_center[0] - 50, frame_center[1] + 50 + i * 25),
                font_scale=0.5, font_thickness=1,
                text_color=(255, 255, 255), bg_color=(100, 100, 100)
            )
    
    def _update_statistics(self):
        """更新统计信息"""
        self.frame_count += 1
        self.fps_counter += 1
        
        current_time = time.time()
        if current_time - self.last_fps_time >= 1.0:
            self.current_fps = self.fps_counter
            self.fps_counter = 0
            self.last_fps_time = current_time
    
    def _cleanup(self):
        """清理资源"""
        self.is_running = False
        self.camera_handler.close_camera()
        cv2.destroyAllWindows()
        print("程序已退出")

def main():
    """主函数"""
    parser = argparse.ArgumentParser(description='摄像头坐姿检测系统 - 精细化评估版本')
    parser.add_argument('--camera', type=int, default=0, help='摄像头ID (默认: 0)')
    parser.add_argument('--debug', action='store_true', help='启用调试模式')
    parser.add_argument('--no-display', action='store_true', help='不显示窗口（仅用于测试）')
    
    args = parser.parse_args()
    
    try:
        app = PostureDetectionApp(camera_id=args.camera, show_debug=args.debug)
        
        if args.no_display:
            print("测试模式：检测系统初始化成功")
            return
        
        app.run()
        
    except KeyboardInterrupt:
        print("\n用户中断程序")
    except Exception as e:
        print(f"程序运行错误: {e}")
        import traceback
        traceback.print_exc()

if __name__ == "__main__":
    main() 